from enum import Enum
import os
import subprocess
import sys
import time

# 定义要执行的脚本字典
class UnitType(Enum):
    Normal = 1
    Sys = 2
    Quit = 3

UnitTests = [
    ('gc', '陀螺仪 - 校准', UnitType.Normal, 'python calibrate_imu.py'),
    ('gs', '陀螺仪 - 实时状态', UnitType.Normal, 'python ../mini_bdx_runtime/mini_bdx_runtime/raw_imu.py'), 
    (), 
    ('f', '足底微动开关', UnitType.Normal, 'python ../mini_bdx_runtime/mini_bdx_runtime/feet_contacts.py'),
    (), 
    ('bt', '手柄 - 启动蓝牙配对程序', UnitType.Normal, 'bluetoothctl'), 
    ('jx', '手柄 - XBox', UnitType.Normal, 'python ../mini_bdx_runtime/mini_bdx_runtime/xbox_controller.py'), 
    ('jp', '手柄 - PS4', UnitType.Normal, 'python ../mini_bdx_runtime/mini_bdx_runtime/ps4_controller.py'), 
    (), 
    ('mpi', '舵机参数 - 批量初始化参数', UnitType.Normal, 'python motor_params.py --init'), 
    ('mpl', '舵机参数 - 读取重要参数', UnitType.Normal, 'python motor_params.py --list'), 
    ('mzg', '舵机逻辑零位 - 进入逻辑零位', UnitType.Normal, 'python motor_logical_zero.py --go_zero'), 
    ('mza', '舵机逻辑零位 - 执行逻辑零位（中位）校准', UnitType.Normal, 'python motor_logical_zero.py --apply_current'), 
    ('mi', '舵机 - 进入初始位', UnitType.Normal, 'python motor_init_pos.py'), 
    ('mo', '舵机 - 关闭扭矩', UnitType.Normal, 'python motor_turn_off.py'),
    ('mv', '舵机 - 查看电压', UnitType.Normal, 'python motor_check_voltage.py'), 
    (), 
    ('w', '开始行走', UnitType.Normal, 'python v2_rl_walk_mujoco.py --onnx_model_path ~/Open_Duck_Mini/BEST_WALK_ONNX_2.onnx'), 
    (), 
    ('r', '重启操作系统', UnitType.Sys, 'sudo reboot'), 
    ('h', '安全关机', UnitType.Sys, 'sudo halt'), 
    ('q', '退出菜单', UnitType.Quit, ''), 
]

def clear_screen():
    """清空控制台屏幕"""
    os.system('cls' if os.name == 'nt' else 'clear')

def display_menu():
    # clear_screen()

    print('\n\n')
    print("========= Open Duck Mini 单元测试 =========")
    print()

    for unit_info in UnitTests:
        if len(unit_info) == 4:
            if unit_info[2] == UnitType.Normal or unit_info[2] == UnitType.Sys:
                print(f"  [{unit_info[0]}] {unit_info[1]}（{unit_info[3]}）")
            elif unit_info[2] == UnitType.Quit:
                print(f"  [{unit_info[0]}] {unit_info[1]}")
        else:
            print('')

    print()

def run_script(cmd_line):
    try:
        # 使用当前Python解释器运行脚本
        print(f"将要执行脚本: {cmd_line}……\n")
        subprocess.run(cmd_line, check=True, shell=True)
    except KeyboardInterrupt:
            print("执行脚本时被用户中断")
    except BaseException as e:
        print(f"错误: 脚本 {cmd_line} 执行失败: {e}")


def find_unit_info(choice):
    for unit_info in UnitTests:
        if len(unit_info) == 4 and unit_info[0] == choice:
            return unit_info
        
    return None


def main():
    last_unit_info = None

    while True:
        display_menu()

        try:
            if last_unit_info is None:
                choice = input("请选择要执行的编号: ").strip().lower()
                unit_info = find_unit_info(choice)
            else:
                choice = input(f"请选择要执行的编号，若直接按回车则继续上次的选择（{last_unit_info[0]}）: ").strip().lower()
                if choice == '':
                    unit_info = last_unit_info
                else:
                    unit_info = find_unit_info(choice)
            
            if unit_info is not None:
                if unit_info[2] == UnitType.Normal:
                    last_unit_info = unit_info
                    run_script(unit_info[3])
                    continue
                elif unit_info[2] == UnitType.Sys:
                    run_script(unit_info[3])
                    break
                elif unit_info[2] == UnitType.Quit:
                    break
            else:
                print("无效的选择，请重新选择。")
                time.sleep(1)
                continue
        except KeyboardInterrupt:
            print("用户中断，程序退出。")
            exit(0)
        except BaseException as e:
            print(f"错误: {e}，程序退出。")
            exit(0)

if __name__ == "__main__":
    main()